use image::{ImageBuffer, Rgb, RgbImage};
use num_complex::Complex64;
use std::fmt::Debug; // 导入Debug trait
use std::io::{self, Write};

// 程序参数结构体
#[derive(Debug)]
struct Args {
    fractal: String,          // 分形类型 (mandelbrot 或 julia)
    width: u32,               // 图像宽度
    height: u32,              // 图像高度
    max_iter: u32,            // 最大迭代次数
    output: String,           // 输出文件路径
    julia_c_real: f64,        // Julia集合的参数c（实部）
    julia_c_imag: f64,        // Julia集合的参数c（虚部）
}

// 从键盘读取输入并解析为指定类型，支持默认值
// 新增 Debug 约束：T: FromStr + Debug
fn read_input<T: std::str::FromStr + Debug>(prompt: &str, default: T) -> T
where
    <T as std::str::FromStr>::Err: std::fmt::Debug,
{
    loop {
        // 现在可以安全使用 {:?} 格式化T类型
        print!("{} [默认: {:?}]: ", prompt, default);
        io::stdout().flush().unwrap();
        
        let mut input = String::new();
        io::stdin().read_line(&mut input).expect("读取输入失败");
        let input = input.trim();
        
        // 空输入使用默认值
        if input.is_empty() {
            return default;
        }
        
        // 尝试解析输入
        match input.parse() {
            Ok(value) => return value,
            Err(_) => println!("输入格式错误，请重新输入！"),
        }
    }
}

// 读取分形类型（确保是mandelbrot或julia）
fn read_fractal_type() -> String {
    loop {
        print!("请输入分形类型 (mandelbrot 或 julia): ");
        io::stdout().flush().unwrap();
        
        let mut input = String::new();
        io::stdin().read_line(&mut input).expect("读取输入失败");
        let fractal = input.trim().to_lowercase();
        
        if fractal == "mandelbrot" || fractal == "julia" {
            return fractal;
        } else {
            println!("不支持的分形类型！请使用 'mandelbrot' 或 'julia'");
        }
    }
}

// 读取所有参数
fn read_args() -> Args {
    println!("=== Rust 分形生成工具 ===");
    
    // 读取分形类型
    let fractal = read_fractal_type();
    
    // 读取通用参数（提供默认值）
    let width = read_input("请输入图像宽度", 800);
    let height = read_input("请输入图像高度", 600);
    let max_iter = read_input("请输入最大迭代次数", 1000);
    let output = read_input("请输入输出文件路径", "fractal.png".to_string());
    
    // 根据分形类型读取特定参数
    let (julia_c_real, julia_c_imag) = if fractal == "julia" {
        (
            read_input("请输入Julia集合参数c的实部", -0.7),
            read_input("请输入Julia集合参数c的虚部", 0.27015),
        )
    } else {
        // 曼德博集合默认参数（无实际作用，仅占位）
        (10.0,10.0)
    };
    
    Args {
        fractal,
        width,
        height,
        max_iter,
        output,
        julia_c_real,
        julia_c_imag,
    }
}

// 曼德博集合计算
fn mandelbrot(c: Complex64, max_iter: u32) -> u32 {
    let mut z = Complex64::new(0.0, 0.0);
    for i in 0..max_iter {
        z = z * z + c;
        if z.norm_sqr() > 4.0 {
            return i;
        }
    }
    max_iter
}

// 朱利亚集合计算
fn julia(z: Complex64, c: Complex64, max_iter: u32) -> u32 {
    let mut current = z;
    for i in 0..max_iter {
        current = current * current + c;
        if current.norm_sqr() > 4.0 {
            return i;
        }
    }
    max_iter
}

// HSV转RGB颜色（保持核心逻辑）
fn hsv_to_rgb(h: f32, s: f32, v: f32) -> (u8, u8, u8) {
    let c = v * s;
    let x = c * (1.0 - ((h / 60.0) % 2.0 - 1.0).abs());
    let m = v - c;

    let (r, g, b) = match (h / 60.0).floor() as i32 {
        0 => (c, x, 0.0),
        1 => (x, c, 0.0),
        2 => (0.0, c, x),
        3 => (0.0, x, c),
        4 => (x, 0.0, c),
        5 => (c, 0.0, x),
        _ => (0.0, 0.0, 0.0),
    };

    (
        ((r + m) * 255.0) as u8,
        ((g + m) * 255.0) as u8,
        ((b + m) * 255.0) as u8,
    )
}

// 分形颜色渲染（优化颜色效果）
fn get_fractal_color(iterations: u32, max_iter: u32) -> (u8, u8, u8) {
    // 集合内的点用深灰色，避免纯黑压抑
    if iterations == max_iter {
        return (20, 20, 30);
    }

    // 高饱和度（0.9）+ 降低明度（0.65），避免过于鲜艳
    let saturation = 0.9;
    let value = 0.65;

    // 阶梯化迭代次数，强化颜色界限（分组映射）
    let step = max_iter / 30; // 分成30组，每组颜色相同
    let stepped_iter = if step == 0 { iterations } else { (iterations / step) * step };

    // 色调范围：180°（青）~ 360°（红），避开刺眼的黄绿
    let hue = 180.0 + (stepped_iter as f32 / max_iter as f32) * 180.0;

    hsv_to_rgb(hue, saturation, value)
}

// 渲染分形图像
fn render_fractal(args: &Args) -> Result<(), image::ImageError> {
    let width = args.width;
    let height = args.height;
    let max_iter = args.max_iter;
    let output_path = &args.output;

    let mut img: RgbImage = ImageBuffer::new(width, height);

    match args.fractal.as_str() {
        "mandelbrot" => {
            // 曼德博集合默认坐标范围（完整视图）
            let x_min = -2.0;
            let x_max = 1.0;
            let y_min = -1.0;
            let y_max = 1.0;

            for (x, y, pixel) in img.enumerate_pixels_mut() {
                let real = x_min + (x as f64 / width as f64) * (x_max - x_min);
                let imag = y_min + (y as f64 / height as f64) * (y_max - y_min);
                let c = Complex64::new(real, imag);

                let iter = mandelbrot(c, max_iter);
                let (r, g, b) = get_fractal_color(iter, max_iter);
                *pixel = Rgb([r, g, b]);
            }
        }
        "julia" => {
            let x_min = -1.5;
            let x_max = 1.5;
            let y_min = -1.0;
            let y_max = 1.0;

            let c = Complex64::new(args.julia_c_real, args.julia_c_imag);

            for (x, y, pixel) in img.enumerate_pixels_mut() {
                let real = x_min + (x as f64 / width as f64) * (x_max - x_min);
                let imag = y_min + (y as f64 / height as f64) * (y_max - y_min);
                let z = Complex64::new(real, imag);

                let iter = julia(z, c, max_iter);
                let (r, g, b) = get_fractal_color(iter, max_iter);
                *pixel = Rgb([r, g, b]);
            }
        }
        _ => panic!("不支持的分形类型！"),
    }

    img.save(output_path)?;
    println!("分形图像已保存至: {}", output_path);
    Ok(())
}

fn main() -> Result<(), Box<dyn std::error::Error>> {
    let args = read_args();
    render_fractal(&args)?;
    Ok(())
}